package com.neusoft.databus.core.loadbalance;

import java.util.Collection;
import java.util.List;
import java.util.Optional;

import org.springframework.util.StringUtils;

import com.netflix.loadbalancer.ILoadBalancer;
import com.netflix.loadbalancer.RoundRobinRule;
import com.netflix.loadbalancer.Server;
import com.netflix.niws.loadbalancer.DiscoveryEnabledServer;
import com.neusoft.bizcore.common.cache.CacheService;
import com.neusoft.bizcore.web.utils.SpringUtil;
import com.neusoft.databus.core.constant.DatabusCoreConstant;

import lombok.extern.slf4j.Slf4j;

@Slf4j
public class StatefulRoundRobinRule extends RoundRobinRule {

    public StatefulRoundRobinRule() {
        super();
    }

    public StatefulRoundRobinRule(ILoadBalancer lb) {
        super();
    }

    @Override
    public Server choose(Object key) {
        final feign.Request request = (feign.Request) key;
        final Collection<String> headers = request.headers().get("stateIdentifier");
        // 如果没有设置头信息，说明服务无需状态
        if (null == headers || headers.isEmpty()) {
            return this.choose(this.getLoadBalancer(), key);
        }
        final String stateIdentifier = headers.iterator().next();
        // 如果头信息为空，说明服务无需状态
        if (StringUtils.isEmpty(stateIdentifier)) {
            return this.choose(this.getLoadBalancer(), key);
        }

        // 以下是带状态服务选择逻辑
        String cachedNode = "";
        final CacheService cacheService = SpringUtil.getBean(CacheService.class);
        if (cacheService.isInCache(DatabusCoreConstant.CACHE_NAME_AGENT_SERVICE_NODE_ROUTE, stateIdentifier)) {
            cachedNode =
                    (String) cacheService.get(DatabusCoreConstant.CACHE_NAME_AGENT_SERVICE_NODE_ROUTE, stateIdentifier);
        }
        Server chosenServer = null;
        if (!StringUtils.isEmpty(cachedNode)) {
            chosenServer = this.chooseWithStateIdentifier(this.getLoadBalancer(), stateIdentifier, cachedNode);
            if (null == chosenServer) {
                chosenServer = this.chooseFirstTime(stateIdentifier, this.getLoadBalancer(), key);
                if (null == chosenServer) {
                    return null;
                }
                if (StatefulRoundRobinRule.log.isDebugEnabled()) {
                    StatefulRoundRobinRule.log.debug("采集网关{} 数据，路由到 {} 节点", stateIdentifier,
                            ((DiscoveryEnabledServer) chosenServer).getInstanceInfo().getId());
                }
                return chosenServer;
            } else {
                if (StatefulRoundRobinRule.log.isDebugEnabled()) {
                    StatefulRoundRobinRule.log.debug("采集网关{}数据，路由到 {} 节点", stateIdentifier, cachedNode);
                }
                return chosenServer;
            }
        } else {

            chosenServer = this.chooseFirstTime(stateIdentifier, this.getLoadBalancer(), key);
            if (null == chosenServer) {
                return null;
            }
            if (StatefulRoundRobinRule.log.isDebugEnabled()) {
                StatefulRoundRobinRule.log.debug("采集网关{} 数据，路由到 {} 节点", stateIdentifier,
                        ((DiscoveryEnabledServer) chosenServer).getInstanceInfo().getId());
            }
            return chosenServer;
        }
    }

    public synchronized Server chooseFirstTime(String stateIdentifier, ILoadBalancer lb, Object key) {
        Server chosenServer = null;
        final CacheService cacheService = SpringUtil.getBean(CacheService.class);
        if (cacheService.isInCache(DatabusCoreConstant.CACHE_NAME_AGENT_SERVICE_NODE_ROUTE, stateIdentifier)) {
            final String cachedNode =
                    (String) cacheService.get(DatabusCoreConstant.CACHE_NAME_AGENT_SERVICE_NODE_ROUTE, stateIdentifier);
            chosenServer = this.chooseWithStateIdentifier(lb, stateIdentifier, cachedNode);
            return chosenServer;
        }

        chosenServer = this.choose(this.getLoadBalancer(), key);
        cacheService.put(DatabusCoreConstant.CACHE_NAME_AGENT_SERVICE_NODE_ROUTE, stateIdentifier,
                ((DiscoveryEnabledServer) chosenServer).getInstanceInfo().getId());
        return chosenServer;
    }

    public Server chooseWithStateIdentifier(ILoadBalancer lb, String stateIdentifier, String cachedAgentNode) {
        final List<Server> servers = lb.getAllServers();

        if (StatefulRoundRobinRule.log.isDebugEnabled()) {
            final StringBuilder sb = new StringBuilder();
            for (final Server server : servers) {
                final String instanceId = ((DiscoveryEnabledServer) server).getInstanceInfo().getId();
                sb.append("|").append(instanceId);

            }
            StatefulRoundRobinRule.log.debug("采集节点集群元信息: {}", sb.toString());
        }
        final Optional<Server> opt = servers.stream().filter(server -> {
            final DiscoveryEnabledServer dserver = (DiscoveryEnabledServer) server;
            final String instanceId = dserver.getInstanceInfo().getId();
            return instanceId.equals(cachedAgentNode);
        }).findFirst();

        if (!opt.isPresent()) {
            final CacheService cacheService = SpringUtil.getBean(CacheService.class);
            cacheService.remove(DatabusCoreConstant.CACHE_NAME_AGENT_SERVICE_NODE_ROUTE, stateIdentifier);
            return null;
        }

        return opt.get();
    }

}
